#include "gc.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkComposeImageFilter.h"
#include "itkImageRegionIterator.h"

typedef itk::ImageFileReader<GrayImage> GrayReader;
typedef itk::ImageFileWriter<RgbImage> RgbWriter;
typedef unsigned char IsoPixel;
typedef itk::Image<IsoPixel, IMAGE_DIMENSION> IsoImage;
typedef itk::ImageFileWriter<IsoImage> IsoWriter;
typedef itk::ComposeImageFilter<GrayImage, RgbImage> RgbComposer;
const int RGB_COLOR_NUM = 6;
const GrayPixel colorTab[RGB_COLOR_NUM][3] = 
  {{0, 0, 32}, {0, 32, 0}, {0, 32, 32}, 
   {32, 0, 32}, {32, 32, 0}, {32, 32, 32}};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
void fillComponent (Component const& component, RgbImage::Pointer image, 
		    RgbPixel const& fgPixel, double bgPortion)
{
  double fgPortion = 1.0 - bgPortion;
  for (ComponentPointSet::const_iterator itr = component.points.begin(); 
       itr != component.points.end(); itr++) {
    RgbImage::IndexType index;
    index[0] = itr->x;
    index[1] = itr->y;
    RgbPixel bgPixel = image->GetPixel(index);
    RgbPixel drawPixel;
    drawPixel.Set(bgPixel.GetRed() * bgPortion + fgPixel.GetRed() * fgPortion, 
		  bgPixel.GetGreen() * bgPortion + fgPixel.GetGreen() * fgPortion, 
		  bgPixel.GetBlue() * bgPortion + fgPixel.GetBlue() * fgPortion);
    image->SetPixel(index, drawPixel);
  }
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
RgbImage::Pointer GrayToRgb (const char* imageFilePath)
{
  GrayReader::Pointer reader = GrayReader::New();
  reader->SetFileName(imageFilePath);
  RgbComposer::Pointer composer = RgbComposer::New();
  composer->SetInput1(reader->GetOutput());
  composer->SetInput2(reader->GetOutput());
  composer->SetInput3(reader->GetOutput());
  composer->Update();
  return composer->GetOutput();
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
void writeRgbImage (RgbImage::Pointer rgbImage, const char* imageFilePath)
{
  RgbWriter::Pointer writer = RgbWriter::New();
  writer->SetInput(rgbImage);
  writer->SetFileName(imageFilePath);
  writer->Update();
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
void drawRGB (Volume const& volume, Graph3D const& graph, 
	      SegmentationGraph fwdSegGraph, 
	      SegmentationGraph bwdSegGraph, 
	      std::vector<std::string> backgroundPaths, 
	      std::vector<std::string> outputPaths, 
	      VerboseLevel verbose)
{
  SegmentationGraph bothSegGraph = fwdSegGraph.getIntersection(bwdSegGraph);
  SegmentationGraph exFwdSegGraph = fwdSegGraph.getDifference(bothSegGraph);
  SegmentationGraph exBwdSegGraph = bwdSegGraph.getDifference(bothSegGraph);
  bothSegGraph.sort();
  exFwdSegGraph.sort();
  exBwdSegGraph.sort();
  RgbPixel fgPixel_both;
  RgbPixel fgPixel_fwd;
  RgbPixel fgPixel_bwd;
  fgPixel_both.Set(0, 255, 0);
  fgPixel_fwd.Set(255, 0, 0);
  fgPixel_bwd.Set(0, 0, 255);
  double bgPortion = 0.2;
  int layerNum = volume.layers.size();
  int currentLayerNo = 0;
  SegGraphSegNodeSet::const_iterator itr_both = bothSegGraph.segNodeSet.begin();
  SegGraphSegNodeSet::const_iterator itr_fwd = exFwdSegGraph.segNodeSet.begin();
  SegGraphSegNodeSet::const_iterator itr_bwd = exBwdSegGraph.segNodeSet.begin();
  for (int layerNo = 0; layerNo < layerNum; layerNo++) {
    RgbImage::Pointer rgbImage = GrayToRgb(backgroundPaths[layerNo].c_str());
    while (itr_both->layerNo == layerNo) {
      int componentNo = graph.nodes[layerNo][itr_both->nodeNo].componentNo;
      fillComponent(volume.layers[layerNo].components[componentNo], rgbImage, 
		    fgPixel_both, bgPortion);
      itr_both++;
    }
    while (itr_fwd->layerNo == layerNo) {
      int componentNo = graph.nodes[layerNo][itr_fwd->nodeNo].componentNo;
      fillComponent(volume.layers[layerNo].components[componentNo], rgbImage, 
		    fgPixel_fwd, bgPortion);
      itr_fwd++;
    }
    while (itr_bwd->layerNo == layerNo) {
      int componentNo = graph.nodes[layerNo][itr_bwd->nodeNo].componentNo;
      fillComponent(volume.layers[layerNo].components[componentNo], rgbImage, 
		    fgPixel_bwd, bgPortion);
      itr_bwd++;
    }
    writeRgbImage(rgbImage, outputPaths[layerNo].c_str());
    if (verbose == VERBOSE_HIGH) {
      std::cout << "Visualization: layer " << layerNo << " done." << std::endl;
    }
  }
}
